import json

import numpy.matlib
import numpy as np
import pandas as pd
from scipy.optimize import minimize

from fof_model import FactorModelHelper, Portfolio


# load target & pool
target = pd.read_csv("./target_portfolio.csv")
stock_pool = pd.read_csv("./all_stock.csv").drop_duplicates().sample(500)

# generate model
fm_helper = FactorModelHelper("./data")
cfg = json.load(open("./model_config.json"))
model = fm_helper.get_model(cfg)


# initialize stock portfolio
stock_portfolio = Portfolio()
for i in range(len(target)):
    rec = target.iloc[i]
    rtn = fm_helper.returns_helper.get_stock_return(rec["stock"])
    stock_portfolio.add_asset(rtn, rec["weight"])
stock_portfolio.analyse_assets(model)
stock_info = stock_portfolio.get_portfolio_exposures()
stock_alpha = stock_info["alpha"]
stock_risk = stock_info["indivial_risk"]
stock_exposure = stock_info["exposure"]
stock_exposure = [stock_exposure[k] for k in model.factor_list]
target_exposure = np.array(stock_exposure)

# initialize fund portfolio
fund_portfolio = Portfolio()
stock_pool_touse = []
for i in range(len(stock_pool)):
    rec = stock_pool.iloc[i]
    rtn = fm_helper.returns_helper.get_stock_return(rec["stockid"])
    if len(rtn) < 300:
        continue
    fund_portfolio.add_asset(rtn)
    stock_pool_touse.append(rec["stockid"])
stock_pool = pd.DataFrame({"stockid": stock_pool_touse})
fund_portfolio.analyse_assets(model)

# minimize portfolio risk
X = fund_portfolio.exposures.values
F = model.covmatrix
D = np.diagflat(fund_portfolio.indivial_risk**2)


def min_func(w):
    w = w.reshape([-1, 1])
    tmp = X @ F @ X.T + D
    tmp = w.T @ tmp @ w
    return tmp[0, 0]


def max_diff(w):
    w = w.reshape([-1, 1])
    current = w.T @ X
    diff = current / target_exposure - 1
    diff = max(abs(diff))
    return diff


def assets(w):
    return sum(w > 0)


constrants = [
    {"type": "ineq", "fun": lambda w: 0.02 - max_diff(w)},
    {"type": "ineq", "fun": lambda w: 1 - sum(w)},
    {"type": "ineq", "fun": lambda w: 15 - assets(w)},
]

bnds = tuple([(0, 1)] * len(stock_pool))

w_init = np.array([1 / len(stock_pool)] * len(stock_pool))
res = minimize(min_func, w_init, bounds=bnds, method="SLSQP", constraints=constrants)

"""
TODO - Be adviced, this problem cannot be solved correctly right now.
However, the exposures are close enough.
"""


print("---------------")
print("fund_id\tweight")
for i in range(len(res.x)):
    cur_weight = round(res.x[i], 4) * 100
    if cur_weight > 0:
        print(f"{stock_pool.iloc[i]['stockid']} \t {round(cur_weight, 2)}%")

print("\n---------------")
print("fund_exp\ttarget\tdiff")
wx = res.x
wx = wx.reshape([-1, 1])
fpe = wx.T @ X
fpe = fpe[0]
fpe
for i in range(len(fpe)):
    print(
        f"{round(fpe[i], 6)} \t {round(target_exposure[i], 6)} \t {round(fpe[i] - target_exposure[i], 6)}"
    )

print("\n---------------")
print("fund_risk\tori_risk")
print(f"{round(np.sqrt(res.fun), 6)}\t{round(stock_risk, 6)}")

print("\n---------------")
alpha = (wx.T @ fund_portfolio.alpha.values)[0, 0]
print("fund_alpha\tori_alpha")
print(f"{round(alpha, 6)}\t{round(stock_alpha, 6)}")
